"""
WebSocket集成测试
================

测试WebSocket服务的完整功能
包括连接、订阅、数据推送等
"""

import pytest
import asyncio
import json
from unittest.mock import AsyncMock, Mock, patch
from datetime import datetime, timezone

from app.services.websocket.service import WebSocketService, MessageType
from app.services.data.push_manager import DataPushManager, DataType


@pytest.mark.asyncio
class TestWebSocketIntegration:
    """WebSocket集成测试类"""
    
    @pytest.fixture
    def websocket_service(self):
        """创建WebSocket服务实例"""
        return WebSocketService(max_connections=10)
    
    @pytest.fixture
    def mock_websocket(self):
        """Mock WebSocket连接"""
        websocket = AsyncMock()
        websocket.send = AsyncMock()
        websocket.recv = AsyncMock()
        websocket.close = AsyncMock()
        websocket.closed = False
        websocket.client_state = Mock()
        websocket.client_state.value = 1  # CONNECTED状态
        return websocket
    
    async def test_websocket_connection_flow(self, websocket_service, mock_websocket):
        """测试WebSocket连接流程"""
        client_id = "test_client_1"
        
        # 测试连接
        success = await websocket_service.connect(mock_websocket, client_id)
        assert success
        assert client_id in websocket_service.connections
        assert websocket_service.get_connection_count() == 1
        
        # 测试订阅
        symbols = ["000001.SZ", "000002.SZ"]
        success = await websocket_service.subscribe(client_id, symbols)
        assert success
        assert symbols[0] in websocket_service.subscriptions
        assert symbols[1] in websocket_service.subscriptions
        
        # 测试数据广播
        test_data = {
            "price": 100.0,
            "change": 1.5,
            "changePercent": 1.52,
            "volume": 1000000
        }
        
        success_count = await websocket_service.broadcast_data(symbols[0], test_data)
        assert success_count == 1  # 应该推送给1个客户端
        
        # 测试取消订阅
        success = await websocket_service.unsubscribe(client_id, [symbols[0]])
        assert success
        assert symbols[0] not in websocket_service.subscriptions
        
        # 测试断开连接
        success = await websocket_service.disconnect(client_id)
        assert success
        assert client_id not in websocket_service.connections
        assert websocket_service.get_connection_count() == 0
    
    async def test_websocket_message_handling(self, websocket_service, mock_websocket):
        """测试WebSocket消息处理"""
        client_id = "test_client_2"
        await websocket_service.connect(mock_websocket, client_id)
        
        # 测试订阅消息
        subscribe_message = {
            "type": "subscribe",
            "data": {"symbols": ["000001.SZ"]}
        }
        success = await websocket_service.handle_message(client_id, json.dumps(subscribe_message))
        assert success
        assert "000001.SZ" in websocket_service.subscriptions
        
        # 测试取消订阅消息
        unsubscribe_message = {
            "type": "unsubscribe", 
            "data": {"symbols": ["000001.SZ"]}
        }
        success = await websocket_service.handle_message(client_id, json.dumps(unsubscribe_message))
        assert success
        assert "000001.SZ" not in websocket_service.subscriptions
        
        # 测试心跳消息
        ping_message = {
            "type": "ping",
            "data": {}
        }
        success = await websocket_service.handle_message(client_id, json.dumps(ping_message))
        assert success
        
        # 测试无效消息
        invalid_message = "invalid json"
        success = await websocket_service.handle_message(client_id, invalid_message)
        assert not success
    
    async def test_websocket_broadcast_performance(self, websocket_service):
        """测试WebSocket广播性能"""
        # 创建多个连接
        connections = []
        for i in range(5):
            websocket = AsyncMock()
            websocket.send = AsyncMock()
            client_id = f"client_{i}"
            await websocket_service.connect(websocket, client_id)
            await websocket_service.subscribe(client_id, ["000001.SZ"])
            connections.append((websocket, client_id))
        
        # 测试广播性能
        test_data = {"price": 100.0, "volume": 1000000}
        
        start_time = datetime.now()
        success_count = await websocket_service.broadcast_data("000001.SZ", test_data)
        end_time = datetime.now()
        
        assert success_count == 5  # 应该推送给5个客户端
        assert (end_time - start_time).total_seconds() < 1.0  # 应该在1秒内完成
        
        # 验证所有连接都收到了消息
        for websocket, _ in connections:
            assert websocket.send.called
    
    async def test_websocket_connection_limits(self, websocket_service):
        """测试WebSocket连接数限制"""
        # 创建最大连接数
        connections = []
        for i in range(websocket_service.max_connections):
            websocket = AsyncMock()
            websocket.send = AsyncMock()
            client_id = f"client_{i}"
            success = await websocket_service.connect(websocket, client_id)
            assert success
            connections.append((websocket, client_id))
        
        # 尝试创建超出限制的连接
        extra_websocket = AsyncMock()
        extra_websocket.close = AsyncMock()
        success = await websocket_service.connect(extra_websocket, "extra_client")
        assert not success
        assert extra_websocket.close.called
    
    async def test_websocket_metrics(self, websocket_service, mock_websocket):
        """测试WebSocket指标收集"""
        client_id = "metrics_test_client"
        
        # 初始指标
        initial_metrics = websocket_service.get_metrics()
        assert initial_metrics["active_connections"] == 0
        assert initial_metrics["total_connections"] == 0
        
        # 连接后指标
        await websocket_service.connect(mock_websocket, client_id)
        metrics = websocket_service.get_metrics()
        assert metrics["active_connections"] == 1
        assert metrics["total_connections"] == 1
        
        # 断开连接后指标
        await websocket_service.disconnect(client_id)
        metrics = websocket_service.get_metrics()
        assert metrics["active_connections"] == 0
        assert metrics["total_connections"] == 1  # 总连接数不会减少
        assert metrics["total_disconnections"] == 1


@pytest.mark.asyncio
class TestDataPushManager:
    """数据推送管理器测试类"""
    
    @pytest.fixture
    def mock_market_data_service(self):
        """Mock市场数据服务"""
        service = AsyncMock()
        service.get_realtime_data = AsyncMock(return_value={
            "price": 100.0,
            "change": 1.5,
            "changePercent": 1.52,
            "volume": 1000000
        })
        return service
    
    @pytest.fixture
    def data_push_manager(self, mock_market_data_service):
        """创建数据推送管理器实例"""
        return DataPushManager(mock_market_data_service)
    
    async def test_data_push_manager_initialization(self, data_push_manager):
        """测试数据推送管理器初始化"""
        assert data_push_manager.is_running == False
        assert len(data_push_manager.push_configs) == 5  # 5种数据类型
        assert DataType.REALTIME_TICK in data_push_manager.push_configs
        assert DataType.INDEX_DATA in data_push_manager.push_configs
    
    async def test_data_push_manager_start_stop(self, data_push_manager):
        """测试数据推送管理器启动和停止"""
        # 启动
        await data_push_manager.start()
        assert data_push_manager.is_running == True
        assert len(data_push_manager.push_tasks) > 0
        
        # 停止
        await data_push_manager.stop()
        assert data_push_manager.is_running == False
        assert len(data_push_manager.push_tasks) == 0
    
    async def test_manual_data_push(self, data_push_manager, mock_market_data_service):
        """测试手动数据推送"""
        symbols = ["000001.SZ", "000002.SZ"]
        
        result = await data_push_manager.push_manual_data(DataType.REALTIME_TICK, symbols)
        
        assert result["data_type"] == "realtime_tick"
        assert result["symbols"] == symbols
        assert result["success_count"] >= 0
        assert result["error_count"] >= 0
        
        # 验证调用了市场数据服务
        assert mock_market_data_service.get_realtime_data.call_count == len(symbols)
    
    def test_push_config_update(self, data_push_manager):
        """测试推送配置更新"""
        from app.services.data.push_manager import PushConfig
        
        new_config = PushConfig(
            data_type=DataType.REALTIME_TICK,
            symbols=["000001.SZ"],
            interval_seconds=2,
            enabled=True
        )
        
        data_push_manager.update_push_config(DataType.REALTIME_TICK, new_config)
        
        config = data_push_manager.push_configs[DataType.REALTIME_TICK]
        assert config.interval_seconds == 2
        assert config.symbols == ["000001.SZ"]
    
    def test_symbol_management(self, data_push_manager):
        """测试股票代码管理"""
        symbols = ["000001.SZ", "000002.SZ", "000003.SZ"]
        
        # 添加股票代码
        data_push_manager.add_symbols(DataType.REALTIME_TICK, symbols)
        config = data_push_manager.push_configs[DataType.REALTIME_TICK]
        assert all(symbol in config.symbols for symbol in symbols)
        
        # 移除股票代码
        data_push_manager.remove_symbols(DataType.REALTIME_TICK, symbols[:2])
        config = data_push_manager.push_configs[DataType.REALTIME_TICK]
        assert "000001.SZ" not in config.symbols
        assert "000002.SZ" not in config.symbols
        assert "000003.SZ" in config.symbols
    
    def test_push_status(self, data_push_manager):
        """测试推送状态获取"""
        status = data_push_manager.get_push_status()
        
        assert "is_running" in status
        assert "push_configs" in status
        assert "active_tasks" in status
        assert "websocket_connections" in status
        assert "websocket_subscriptions" in status
        
        assert status["is_running"] == False
        assert len(status["push_configs"]) == 5


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
